#!/usr/bin/env python3

import os
import mmap
import struct
import time
import numpy as np
import cv2
import rospy
import threading
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
from collections import deque

class TimingStats:
    """计时统计类"""
    
    def __init__(self, max_samples=100):
        self.max_samples = max_samples
        self.stats = {
            'bridge_conversion': deque(maxlen=max_samples),
            'memory_write': deque(maxlen=max_samples),
            'total_process': deque(maxlen=max_samples),
            'callback_time': deque(maxlen=max_samples)
        }
        self.counters = {
            'total_attempts': 0,
            'successful_writes': 0,
            'failed_writes': 0,
            'color_callbacks': 0,
            'depth_callbacks': 0,
            'pairs_processed': 0
        }
        self.lock = threading.Lock()
        self.start_time = time.time()
    
    def add_timing(self, category, duration_ms):
        """添加计时数据"""
        with self.lock:
            if category in self.stats:
                self.stats[category].append(duration_ms)
    
    def increment_counter(self, counter_name):
        """增加计数器"""
        with self.lock:
            if counter_name in self.counters:
                self.counters[counter_name] += 1
    
    def get_stats_summary(self):
        """获取统计摘要"""
        with self.lock:
            summary = {}
            
            # 计算各项时间统计
            for category, times in self.stats.items():
                if times:
                    summary[category] = {
                        'avg': sum(times) / len(times),
                        'max': max(times),
                        'min': min(times),
                        'latest': times[-1] if times else 0,
                        'count': len(times)
                    }
            
            # 添加计数器统计
            summary['counters'] = self.counters.copy()
            
            # 计算运行时间和FPS
            uptime = time.time() - self.start_time
            summary['uptime_seconds'] = uptime
            
            if uptime > 0:
                summary['fps'] = {
                    'color_input': self.counters['color_callbacks'] / uptime,
                    'depth_input': self.counters['depth_callbacks'] / uptime,
                    'processing': self.counters['pairs_processed'] / uptime,
                    'success_rate': (self.counters['successful_writes'] / 
                                   max(1, self.counters['total_attempts']) * 100)
                }
            
            return summary
    
    def reset_stats(self):
        """重置统计数据"""
        with self.lock:
            for category in self.stats:
                self.stats[category].clear()
            for counter in self.counters:
                self.counters[counter] = 0
            self.start_time = time.time()

class SharedMemoryManager:
    """优化的共享内存管理器 - 直接存储原始图像数据，无编码"""
    
    def __init__(self, shm_name="/image_shm", max_size=150 * 1024 * 1024):  # 增加到150MB
        self.shm_name = shm_name
        self.max_size = max_size
        self.lock = threading.Lock()
        self._create_shared_memory()
        self.bridge = CvBridge()
        
        # 添加计时统计
        self.timing_stats = TimingStats()
        
        rospy.loginfo(f"SharedMemoryManager initialized (RAW MODE) - Max size: {max_size/1024/1024:.1f}MB")
    
    def _create_shared_memory(self):
        """创建共享内存区域"""
        try:
            os.unlink(self.shm_name)
        except FileNotFoundError:
            pass
        
        # 创建共享内存文件
        self.fd = os.open(self.shm_name, os.O_CREAT | os.O_RDWR, 0o666)
        os.ftruncate(self.fd, self.max_size)
        
        # 映射共享内存为可读写
        self.shm = mmap.mmap(self.fd, self.max_size, mmap.MAP_SHARED, mmap.PROT_READ | mmap.PROT_WRITE)
    
    def _dtype_to_code(self, dtype):
        """numpy数据类型转换为数字代码"""
        dtype_map = {
            np.uint8: 1,
            np.uint16: 2,
            np.int16: 3,
            np.float32: 4,
            np.float64: 5
        }
        return dtype_map.get(dtype.type, 2)  # 默认uint16
    
    def _code_to_dtype(self, code):
        """数字代码转换为numpy数据类型"""
        code_map = {
            1: np.uint8,
            2: np.uint16,
            3: np.int16,
            4: np.float32,
            5: np.float64
        }
        return code_map.get(code, np.uint16)  # 默认为uint16（常见深度图格式）
    
    def write_images(self, color_msg, depth_msg):
        """将图像写入共享内存 - 直接存储原始数据，无编码"""
        total_start = time.time()
        self.timing_stats.increment_counter('total_attempts')
        
        with self.lock:
            try:
                # 步骤1: ROS消息转换为OpenCV图像
                bridge_start = time.time()
                color_img = self.bridge.imgmsg_to_cv2(color_msg, "bgr8")
                depth_img = self.bridge.imgmsg_to_cv2(depth_msg, "passthrough")
                bridge_time = (time.time() - bridge_start) * 1000
                self.timing_stats.add_timing('bridge_conversion', bridge_time)
                
                rospy.logdebug(f"📸 图像转换完成 - 耗时: {bridge_time:.2f}ms, "
                             f"彩色: {color_img.shape}, 深度: {depth_img.shape}")
                
                # 步骤2: 获取图像属性和数据大小
                color_height, color_width, color_channels = color_img.shape
                depth_height, depth_width = depth_img.shape
                depth_dtype = depth_img.dtype
                
                color_data_size = color_img.nbytes
                depth_data_size = depth_img.nbytes
                
                rospy.logdebug(f"📦 原始数据大小 - 彩色: {color_data_size/1024:.1f}KB, "
                             f"深度: {depth_data_size/1024:.1f}KB")
                
                # 步骤3: 计算总数据大小并检查
                # 数据头结构 (49字节)：
                # - 标志 (5字节): 'VALID'
                # - 时间戳 (8字节): uint64
                # - 彩色图像信息 (12字节): height(4) + width(4) + channels(4)
                # - 深度图像信息 (12字节): height(4) + width(4) + dtype_code(4)
                # - 数据大小 (8字节): color_size(4) + depth_size(4)
                # - 保留字段 (4字节): 用于未来扩展
                header_size = 49
                total_size = header_size + color_data_size + depth_data_size
                
                rospy.logdebug(f"📊 总数据大小: {total_size/1024:.1f}KB (含{header_size}字节头)")
                
                if total_size > self.max_size:
                    rospy.logerr(f"❌ 图像数据太大 - {total_size/1024:.1f}KB > {self.max_size/1024:.1f}KB")
                    self.timing_stats.increment_counter('failed_writes')
                    return False
                
                # 步骤4: 写入共享内存（直接写入原始数据）
                memory_write_start = time.time()
                timestamp = int(time.time() * 1000)
                
                self.shm.seek(0)
                
                # 写入数据头
                self.shm.write(b'VALID')  # 5字节标志
                self.shm.write(struct.pack('Q', timestamp))  # 8字节时间戳
                
                # 彩色图像信息 (12字节)
                self.shm.write(struct.pack('III', color_height, color_width, color_channels))
                
                # 深度图像信息 (12字节) - dtype转换为数字代码
                dtype_code = self._dtype_to_code(depth_dtype)
                self.shm.write(struct.pack('III', depth_height, depth_width, dtype_code))
                
                # 数据大小 (8字节)
                self.shm.write(struct.pack('II', color_data_size, depth_data_size))
                
                # 保留字段 (4字节) - 用于未来扩展
                self.shm.write(struct.pack('I', 0))
                
                # 直接写入原始图像数据（连续的内存块，无压缩）
                self.shm.write(color_img.tobytes())
                self.shm.write(depth_img.tobytes())
                
                memory_write_time = (time.time() - memory_write_start) * 1000
                self.timing_stats.add_timing('memory_write', memory_write_time)
                
                # 步骤5: 计算总时间
                total_time = (time.time() - total_start) * 1000
                self.timing_stats.add_timing('total_process', total_time)
                self.timing_stats.increment_counter('successful_writes')
                
                # 详细性能日志
                rospy.logdebug(f"💾 共享内存写入完成 - 耗时: {memory_write_time:.2f}ms")
                rospy.loginfo_throttle(1, 
                    f"✅ 原始图像处理完成 - 总耗时: {total_time:.2f}ms "
                    f"[转换: {bridge_time:.1f}ms, 写入: {memory_write_time:.1f}ms] "
                    f"(无编码延迟)")
                
                # 性能警告 - 调整阈值，因为去除了编码
                if total_time > 50:  # 降低警告阈值到50ms
                    rospy.logwarn(f"⚠️  处理耗时过长: {total_time:.1f}ms")
                    if bridge_time > 30:
                        rospy.logwarn(f"  - 消息转换慢: {bridge_time:.1f}ms")
                    if memory_write_time > 20:
                        rospy.logwarn(f"  - 内存写入慢: {memory_write_time:.1f}ms")
                
                return True
                
            except Exception as e:
                total_time = (time.time() - total_start) * 1000
                rospy.logerr(f"❌ 写入共享内存错误 - 耗时: {total_time:.2f}ms, 错误: {e}")
                self.timing_stats.increment_counter('failed_writes')
                return False
    
    def read_images(self):
        """从共享内存读取图像 - 直接读取原始数据，无解码"""
        with self.lock:
            try:
                read_start = time.time()
                
                self.shm.seek(0)
                
                # 读取标志
                flag = self.shm.read(5)
                if flag != b'VALID':
                    return None, None, None, None
                
                # 读取时间戳
                timestamp = struct.unpack('Q', self.shm.read(8))[0]
                readable_timestamp = time.strftime('%Y-%m-%d %H:%M:%S', 
                                                 time.localtime(timestamp / 1000))
                
                # 读取彩色图像信息
                color_height, color_width, color_channels = struct.unpack('III', self.shm.read(12))
                
                # 读取深度图像信息
                depth_height, depth_width, dtype_code = struct.unpack('III', self.shm.read(12))
                depth_dtype = self._code_to_dtype(dtype_code)
                
                # 读取数据大小
                color_data_size, depth_data_size = struct.unpack('II', self.shm.read(8))
                
                # 跳过保留字段
                self.shm.read(4)
                
                # 直接读取原始图像数据
                color_bytes = self.shm.read(color_data_size)
                depth_bytes = self.shm.read(depth_data_size)
                
                # 重构图像数组（直接从字节数据重建，无需解码）
                color_img = np.frombuffer(color_bytes, dtype=np.uint8).reshape(
                    (color_height, color_width, color_channels))
                depth_img = np.frombuffer(depth_bytes, dtype=depth_dtype).reshape(
                    (depth_height, depth_width))
                
                read_time = (time.time() - read_start) * 1000
                
                rospy.logdebug(f"📖 原始图像读取完成 - 耗时: {read_time:.2f}ms "
                             f"(无解码延迟)")
                
                return readable_timestamp, str(timestamp), color_img, depth_img
                
            except Exception as e:
                rospy.logerr(f"❌ 读取共享内存错误: {e}")
                return None, None, None, None
    
    def get_performance_report(self):
        """获取性能报告"""
        return self.timing_stats.get_stats_summary()
    
    def reset_performance_stats(self):
        """重置性能统计"""
        self.timing_stats.reset_stats()
        rospy.loginfo("📊 性能统计已重置")
    
    def close(self):
        """关闭共享内存"""
        try:
            # 输出最终统计
            final_stats = self.get_performance_report()
            rospy.loginfo("📈 最终性能统计 (RAW模式):")
            if 'total_process' in final_stats:
                avg_time = final_stats['total_process']['avg']
                max_time = final_stats['total_process']['max']
                rospy.loginfo(f"  总处理时间 - 平均: {avg_time:.1f}ms, 最大: {max_time:.1f}ms")
            
            if 'counters' in final_stats:
                counters = final_stats['counters']
                rospy.loginfo(f"  处理统计 - 成功: {counters['successful_writes']}, "
                            f"失败: {counters['failed_writes']}, "
                            f"总计: {counters['total_attempts']}")
            
            if 'fps' in final_stats:
                fps = final_stats['fps']
                rospy.loginfo(f"  FPS统计 - 处理: {fps['processing']:.1f}, "
                            f"成功率: {fps['success_rate']:.1f}%")
            
            self.shm.close()
            os.close(self.fd)
            os.unlink(self.shm_name)
            rospy.loginfo("🔒 共享内存已关闭")
        except Exception as e:
            rospy.logerr(f"关闭共享内存错误: {e}")

class ImageAcquisition:
    """优化的图像采集类 - 无编码版本"""
    
    def __init__(self):
        self.bridge = CvBridge()
        self.shm_manager = SharedMemoryManager()
        self.color_img = None
        self.depth_img = None
        self.lock = threading.Lock()
        
        # 回调计时
        self.callback_stats = TimingStats(max_samples=50)
        
        # 启动性能监控线程
        self.monitor_thread = threading.Thread(target=self._performance_monitor, daemon=True)
        self.monitor_thread.start()
        
        # ROS订阅
        rospy.Subscriber("/camera/color/image_raw", Image, self.color_callback, queue_size=10)
        rospy.Subscriber("/camera/depth/image_raw", Image, self.depth_callback, queue_size=10)
        
        rospy.loginfo("🚀 ImageAcquisition initialized - RAW MODE (No Encoding)!")
        
    def color_callback(self, msg):
        """彩色图像回调 - 带计时"""
        callback_start = time.time()
        try:
            with self.lock:
                self.color_img = msg
                self.shm_manager.timing_stats.increment_counter('color_callbacks')
                self.try_process_images()
                
            callback_time = (time.time() - callback_start) * 1000
            self.callback_stats.add_timing('callback_time', callback_time)
            
            rospy.logdebug(f"🟢 彩色回调完成 - 耗时: {callback_time:.2f}ms")
            
        except Exception as e:
            callback_time = (time.time() - callback_start) * 1000
            rospy.logerr(f"❌ 彩色回调错误 - 耗时: {callback_time:.2f}ms, 错误: {e}")

    def depth_callback(self, msg):
        """深度图像回调 - 带计时"""
        callback_start = time.time()
        try:
            with self.lock:
                self.depth_img = msg
                self.shm_manager.timing_stats.increment_counter('depth_callbacks')
                self.try_process_images()
                
            callback_time = (time.time() - callback_start) * 1000
            self.callback_stats.add_timing('callback_time', callback_time)
            
            rospy.logdebug(f"🔵 深度回调完成 - 耗时: {callback_time:.2f}ms")
            
        except Exception as e:
            callback_time = (time.time() - callback_start) * 1000
            rospy.logerr(f"❌ 深度回调错误 - 耗时: {callback_time:.2f}ms, 错误: {e}")

    def try_process_images(self):
        """尝试处理图像对"""
        if self.color_img is not None and self.depth_img is not None:
            process_start = time.time()
            
            # 时间戳检查
            color_time = self.color_img.header.stamp.to_sec()
            depth_time = self.depth_img.header.stamp.to_sec()
            time_diff = abs(color_time - depth_time)
            
            rospy.logdebug(f"🔗 图像对匹配 - 时间差: {time_diff*1000:.1f}ms")
            
            if time_diff > 0.1:  # 100ms
                rospy.logwarn_throttle(2, f"⚠️  图像时间戳不匹配: {time_diff*1000:.1f}ms")
            
            # 写入共享内存（原始数据，无编码）
            success = self.shm_manager.write_images(self.color_img, self.depth_img)
            
            process_time = (time.time() - process_start) * 1000
            
            if success:
                self.shm_manager.timing_stats.increment_counter('pairs_processed')
                rospy.loginfo_throttle(3, f"✅ 图像对处理成功 - 总耗时: {process_time:.1f}ms (RAW模式)")
            else:
                rospy.logwarn_throttle(2, f"❌ 图像对处理失败 - 耗时: {process_time:.1f}ms")
            
            # 重置
            self.color_img = None
            self.depth_img = None

    def _performance_monitor(self):
        """性能监控线程"""
        rospy.loginfo("📊 性能监控线程已启动 (RAW模式)")
        
        while not rospy.is_shutdown():
            time.sleep(10)  # 每10秒输出一次详细统计
            
            try:
                # 获取共享内存性能统计
                shm_stats = self.shm_manager.get_performance_report()
                
                if 'counters' in shm_stats and shm_stats['counters']['total_attempts'] > 0:
                    rospy.loginfo("=" * 60)
                    rospy.loginfo("📈 RAW模式性能报告")
                    rospy.loginfo("=" * 60)
                    
                    # 基本统计
                    counters = shm_stats['counters']
                    rospy.loginfo(f"📊 处理统计:")
                    rospy.loginfo(f"  - 总尝试: {counters['total_attempts']}")
                    rospy.loginfo(f"  - 成功: {counters['successful_writes']}")
                    rospy.loginfo(f"  - 失败: {counters['failed_writes']}")
                    rospy.loginfo(f"  - 彩色回调: {counters['color_callbacks']}")
                    rospy.loginfo(f"  - 深度回调: {counters['depth_callbacks']}")
                    
                    # FPS统计
                    if 'fps' in shm_stats:
                        fps = shm_stats['fps']
                        rospy.loginfo(f"🚀 FPS统计:")
                        rospy.loginfo(f"  - 彩色输入: {fps['color_input']:.1f} fps")
                        rospy.loginfo(f"  - 深度输入: {fps['depth_input']:.1f} fps")
                        rospy.loginfo(f"  - 处理输出: {fps['processing']:.1f} fps")
                        rospy.loginfo(f"  - 成功率: {fps['success_rate']:.1f}%")
                    
                    # 详细时间统计 (无编码项)
                    rospy.loginfo(f"⏱️  时间分析 (ms) - RAW模式:")
                    time_categories = [
                        ('bridge_conversion', '消息转换'),
                        ('memory_write', '内存写入'),
                        ('total_process', '总处理')
                    ]
                    
                    for category, display_name in time_categories:
                        if category in shm_stats:
                            stats = shm_stats[category]
                            rospy.loginfo(f"  - {display_name:8}: "
                                        f"平均 {stats['avg']:6.1f}, "
                                        f"最大 {stats['max']:6.1f}, "
                                        f"最小 {stats['min']:6.1f}, "
                                        f"最新 {stats['latest']:6.1f}")
                    
                    # 性能建议
                    rospy.loginfo(f"💡 性能建议:")
                    if 'total_process' in shm_stats:
                        avg_total = shm_stats['total_process']['avg']
                        if avg_total > 50:  # RAW模式阈值更低
                            rospy.logwarn(f"  ⚠️  总处理时间过长 ({avg_total:.1f}ms)")
                            if 'bridge_conversion' in shm_stats and shm_stats['bridge_conversion']['avg'] > 30:
                                rospy.logwarn(f"  - 考虑优化消息转换")
                        else:
                            rospy.loginfo(f"  ✅ 处理速度优秀 ({avg_total:.1f}ms) - RAW模式优势明显!")
                    
                    if 'fps' in shm_stats:
                        processing_fps = shm_stats['fps']['processing']
                        if processing_fps < 10.0:
                            rospy.logwarn(f"  ⚠️  处理FPS偏低 ({processing_fps:.1f})")
                        elif processing_fps > 20.0:
                            rospy.loginfo(f"  ✅ 处理FPS优秀 ({processing_fps:.1f}) - RAW模式高效!")
                    
                    rospy.loginfo("=" * 60)
                    
            except Exception as e:
                rospy.logerr(f"性能监控错误: {e}")

    def get_timing_summary(self):
        """获取计时摘要"""
        shm_stats = self.shm_manager.get_performance_report()
        callback_stats = self.callback_stats.get_stats_summary()
        
        return {
            'shared_memory': shm_stats,
            'callbacks': callback_stats
        }

    def reset_all_stats(self):
        """重置所有统计"""
        self.shm_manager.reset_performance_stats()
        self.callback_stats.reset_stats()
        rospy.loginfo("🔄 所有性能统计已重置")

    def shutdown(self):
        """关闭时清理"""
        rospy.loginfo("🛑 开始关闭ImageAcquisition (RAW模式)...")
        
        # 输出最终摘要
        final_summary = self.get_timing_summary()
        rospy.loginfo("📋 最终计时摘要:")
        rospy.loginfo(f"  运行时间: {final_summary['shared_memory'].get('uptime_seconds', 0):.1f}秒")
        
        self.shm_manager.close()
        rospy.loginfo("✅ ImageAcquisition已关闭")

if __name__ == "__main__":
    rospy.init_node("raw_image_acquisition")
    
    # 设置日志级别以查看详细计时信息
    if rospy.get_param('~debug_timing', False):
        rospy.loginfo("🔍 启用调试计时模式")
        # 可以通过rosparam设置日志级别
    
    acquisition = ImageAcquisition()
    rospy.on_shutdown(acquisition.shutdown)
    
    rospy.loginfo("🎯 RAW模式图像采集节点已启动 - 无编码延迟!")
    rospy.loginfo("💡 提示: 设置 ~debug_timing:=true 查看详细计时信息")
    rospy.loginfo("🚀 性能提升: 去除PNG编码，预计速度提升10-20倍!")
    rospy.spin()